package com.xiao.xin;

import cn.hutool.crypto.digest.BCrypt;
import com.alibaba.fastjson2.JSON;
import com.xiao.xin.comfyui.client.api.ComfyuiApi;
import com.xiao.xin.core.domain.ComfyuiModel;
import com.xiao.xin.core.domain.ComfyuiRequestDto;
import com.xiao.xin.core.domain.User;
import com.xiao.xin.core.service.FreemarkerService;
import com.xiao.xin.core.service.UserService;
import com.xiao.xin.ollama.client.api.OllamaApi;
import com.xiao.xin.ollama.client.dto.OllamaChatDto;
import com.xiao.xin.ollama.client.vo.OllamaChatVo;
import com.xiao.xin.ollama.client.domain.OllamaMessage;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import retrofit2.Call;
import retrofit2.Response;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;

@SpringBootTest
public class test {

    @Autowired
    private ComfyuiApi comfyuiApi;

    @Test
    public void test() throws IOException {
        Call<HashMap> histor = comfyuiApi.getHistoryTasks(2);
        Response<HashMap> response = histor.execute();
        System.out.println(JSON.toJSONString(response.body()));
    }

    @Test
    public void test2() throws IOException {
        Call<HashMap> systemStats = comfyuiApi.getSystemStats();
        Response<HashMap> response = systemStats.execute();
        System.out.println(JSON.toJSONString(response.body()));
    }

    @Test
    public void test3() throws IOException {
        ComfyuiRequestDto queryTask = new ComfyuiRequestDto();
        queryTask.setClientId("6545464");
        queryTask.setPrompt(JSON.parseObject("{\n" + "  \"10\": {\n" + "    \"inputs\": {\n" + "      \"ckpt_name\": \"majicmixRealistic_v7.safetensors\"\n" + "    },\n" + "    \"class_type\": \"CheckpointLoaderSimple\",\n" + "    \"_meta\": {\n" + "      \"title\": \"Load Checkpoint\"\n" + "    }\n" + "  },\n" + "  \"12\": {\n" + "    \"inputs\": {\n" + "      \"text\": \"Landscape\",\n" + "      \"clip\": [\n" + "        \"10\",\n" + "        1\n" + "      ]\n" + "    },\n" + "    \"class_type\": \"CLIPTextEncode\",\n" + "    \"_meta\": {\n" + "      \"title\": \"CLIP Text Encode (Prompt)\"\n" + "    }\n" + "  },\n" + "  \"13\": {\n" + "    \"inputs\": {\n" + "      \"text\": \"Landscape\",\n" + "      \"clip\": [\n" + "        \"10\",\n" + "        1\n" + "      ]\n" + "    },\n" + "    \"class_type\": \"CLIPTextEncode\",\n" + "    \"_meta\": {\n" + "      \"title\": \"CLIP Text Encode (Prompt)\"\n" + "    }\n" + "  },\n" + "  \"14\": {\n" + "    \"inputs\": {\n" + "      \"seed\": 1113283777186661,\n" + "      \"steps\": 20,\n" + "      \"cfg\": 8,\n" + "      \"sampler_name\": \"euler\",\n" + "      \"scheduler\": \"normal\",\n" + "      \"denoise\": 1,\n" + "      \"model\": [\n" + "        \"10\",\n" + "        0\n" + "      ],\n" + "      \"positive\": [\n" + "        \"12\",\n" + "        0\n" + "      ],\n" + "      \"negative\": [\n" + "        \"13\",\n" + "        0\n" + "      ],\n" + "      \"latent_image\": [\n" + "        \"17\",\n" + "        0\n" + "      ]\n" + "    },\n" + "    \"class_type\": \"KSampler\",\n" + "    \"_meta\": {\n" + "      \"title\": \"KSampler\"\n" + "    }\n" + "  },\n" + "  \"15\": {\n" + "    \"inputs\": {\n" + "      \"filename_prefix\": \"ComfyUI\",\n" + "      \"images\": [\n" + "        \"19\",\n" + "        0\n" + "      ]\n" + "    },\n" + "    \"class_type\": \"SaveImage\",\n" + "    \"_meta\": {\n" + "      \"title\": \"Save Image\"\n" + "    }\n" + "  },\n" + "  \"17\": {\n" + "    \"inputs\": {\n" + "      \"width\": 512,\n" + "      \"height\": 512,\n" + "      \"batch_size\": 1\n" + "    },\n" + "    \"class_type\": \"EmptyLatentImage\",\n" + "    \"_meta\": {\n" + "      \"title\": \"Empty Latent Image\"\n" + "    }\n" + "  },\n" + "  \"18\": {\n" + "    \"inputs\": {\n" + "      \"samples\": [\n" + "        \"14\",\n" + "        0\n" + "      ],\n" + "      \"vae\": [\n" + "        \"10\",\n" + "        2\n" + "      ]\n" + "    },\n" + "    \"class_type\": \"VAEDecode\",\n" + "    \"_meta\": {\n" + "      \"title\": \"VAE Decode\"\n" + "    }\n" + "  },\n" + "  \"19\": {\n" + "    \"inputs\": {\n" + "      \"upscale_model\": [\n" + "        \"21\",\n" + "        0\n" + "      ],\n" + "      \"image\": [\n" + "        \"18\",\n" + "        0\n" + "      ]\n" + "    },\n" + "    \"class_type\": \"ImageUpscaleWithModel\",\n" + "    \"_meta\": {\n" + "      \"title\": \"Upscale Image (using Model)\"\n" + "    }\n" + "  },\n" + "  \"21\": {\n" + "    \"inputs\": {\n" + "      \"model_name\": \"通用BSRGANx4.pth\"\n" + "    },\n" + "    \"class_type\": \"UpscaleModelLoader\",\n" + "    \"_meta\": {\n" + "      \"title\": \"Load Upscale Model\"\n" + "    }\n" + "  }\n" + "}"));
        Call<HashMap> systemStats = comfyuiApi.addQueueTask(queryTask);
        Response<HashMap> response = systemStats.execute();
        System.out.println(JSON.toJSONString(response.body()));
    }

    @Autowired
    private UserService userService;

    @Test
    public void test4() {
        User user = new User();
        user.setId(1L);
        user.setMobile("15210307233");
        user.setPassword(BCrypt.hashpw("admin"));
        user.setUsername("admin");
        user.setStatus(1);
        user.setGender(1);
        user.setEmail("admin@126.com");
        user.setNickname("Admin");
        user.setVipLevel(1);
        userService.save(user);
    }

    @Autowired
    private OllamaApi ollamaApi;

    @Test
    public void testChat() throws IOException {
        OllamaChatDto body = new OllamaChatDto();
        body.setModel("qwen2:0.5b");
        OllamaMessage message = new OllamaMessage();
        message.setRole("user");
        message.setContent("帮我把以下内容翻译成英文:美女，风景?");
        body.setMessages(List.of(message));
        Call<OllamaChatVo> chat = ollamaApi.chat(body);
        Response<OllamaChatVo> result = chat.execute();
        System.out.println(result.body().getMessage().getContent());
    }

    @Autowired
    FreemarkerService freemarkerService;

    @Test
    public void test6() throws Exception {
        ComfyuiModel model = new ComfyuiModel();
        model.setPropmt("a photo of an astronaut riding a horse on mars");
        model.setReverse("bad legs");
        model.setSamplerName("Euler");
        model.setScheduler("normal");
        model.setSize(1);
        model.setCfg(7);
        model.setWidth(512);
        model.setHeight(512);
        model.setStep(20);
        model.setSeed(0);
        model.setModelName("majicmixRealistic_v7.safetensors");
        String json = freemarkerService.renderText2Image(model);
        System.out.println(json);
    }
}
